import matplotlib.pyplot as plt
import MDAnalysis as mda
import numpy as np


sel_dict = {'phe': 'resname PHE and not backbone and not (name H* or name CB* or name O*)',
    'tyr': 'resname TYR and not backbone and not (name H* or name CB* or name O*) or name OH',
    'his': 'resname HIS and not backbone and not (name H* or name CB* or name O*)',
    'trp': 'resname TRP and not backbone and not (name H* or name CB* or name O*)',
    'pro': 'resname PRO and not backbone and not (name H* or name O*)',
    'cpd1': 'resname FC and (name C1 or name C2 or name C3 or name C4 or name C5)',
    'cpd2': 'resname FC and (name C6 or name C7 or name C8 or name C9 or name C10)',
}

atnum_per_ring_dict = {'phe': 6,
    'tyr': 7,
    'his': 5,
    'trp': 9,
    'pro': 5,
    'cpd1': 5,
    'cpd2': 5,
}


def get_aminoacid_sidechain_ring_center(ag, aa):
    '''获取原子组的侧链环中心
    参数: 所有原子组, 氨基酸名称
    返回值: 侧链环中心坐标(n×3)
    '''
    ring_atoms = get_aminoacid_sidechain_ring_atoms(ag, aa)
    centers = get_center_from_sidechain_positions(ring_atoms.positions, atnum_per_ring_dict[aa])
    return centers


def get_aminoacid_sidechain_ring_atoms(ag, aa=''):
    '''获取原子组的成环原子(用来计算质心距离)
    参数: 所有原子组, 氨基酸名称
    返回值: 所选原子组
    '''
    return ag.select_atoms(sel_dict[aa])


def get_target_atoms(ag, sel='all'):
    '''获取原子组中要计算距离的目标原子(用来计算原子距离)
    参数: 原子组, 目标原子选择语句
    返回值: 所选原子组
    '''
    return ag.select_atoms(sel)


def get_center_from_sidechain_positions(positions, atnum_per_ring):
    '''根据成环原子坐标计算环中心位置
    参数: 所有成环原子坐标, 每个环包含的原子数
    返回值: 环中心坐标(n×3)
    '''
    return positions.reshape([-1, atnum_per_ring, 3]).mean(axis=1)


def calc_inter_hydrophobic_contacts(u, agA, agB, aa_nameA, aa_nameB, fname, threshold=6):
    '''计算两个原子组间的环间两两接触数
    参数: 原子组A, 原子组B, 氨基酸名A, 氨基酸名B, 输出文件名, 阈值(默认6A)
    '''
    ring_atomsA = get_aminoacid_sidechain_ring_atoms(agA, aa_nameA)
    ring_atomsB = get_aminoacid_sidechain_ring_atoms(agB, aa_nameB)
    if set(ring_atomsA) & set(ring_atomsB):
        raise Exception(aa_nameA + ' in group A and ' + aa_nameB + ' in group B containing same atoms:\n' + str(set(ring_atomsA) & set(ring_atomsB)))
    else:
        with open(fname, 'w') as f:
            f.write('Time (ps)' + ', Contacts Number of ' + aa_nameA + ' & ' + aa_nameB + '\n')
            for ts in u.trajectory:
                centersA = get_center_from_sidechain_positions(ring_atomsA.positions, atnum_per_ring_dict[aa_nameA])
                centersB = get_center_from_sidechain_positions(ring_atomsB.positions, atnum_per_ring_dict[aa_nameB])
                dist = np.array([np.sqrt(np.sum((x - y) ** 2)) for x in centersA for y in centersB])
                f.write(str(u.trajectory.time) + ',' + str(len(dist[dist <= threshold])) + '\n')
                print('Finished:', '{:.2f}'.format(ts.frame / len(u.trajectory) * 100), '%')


def calc_inter_contacts(u, agA, agB, selA, selB, fname, threshold=3.2):
    '''计算两个原子组间的原子间两两接触数
    参数: 原子组A, 原子组B, 氨基酸名A, 氨基酸名B, 输出文件名, 阈值(默认3.2A)
    '''
    ring_atomsA = get_target_atoms(agA, selA)
    ring_atomsB = get_target_atoms(agB, selB)
    if set(ring_atomsA) & set(ring_atomsB):
        raise Exception('group A and group B containing same atoms:\n' + str(set(ring_atomsA) & set(ring_atomsB)))
    else:
        with open(fname, 'w') as f:
            f.write('Time (ps), Contacts Number of group A & group B\n')
            for ts in u.trajectory:
                dist = np.array([np.sqrt(np.sum((x - y) ** 2)) for x in ring_atomsA.positions for y in ring_atomsB.positions])
                f.write(str(u.trajectory.time) + ',' + str(len(dist[dist <= threshold])) + '\n')
                print('Finished:', '{:.2f}'.format(ts.frame / len(u.trajectory) * 100), '%')


def calc_intra_hydrophobic_contacts(u, aa_name, fname, threshold=6):
    '''计算原子组内部环间两两接触数
    参数: 整体原子组, 氨基酸名, 输出文件名, 阈值(默认6A)
    '''
    ring_atoms = get_aminoacid_sidechain_ring_atoms(u, aa_name)
    with open(fname, 'w') as f:
        f.write('Time (ps)' + ',' + 'Contacts Number of ' + aa_name + '\n')
        for ts in u.trajectory:
            centers = get_center_from_sidechain_positions(ring_atoms.positions, atnum_per_ring_dict[aa_name])
            dist = []
            for i, x in enumerate(centers):
                for y in centers[i + 1:]:
                    dist.append(np.sqrt(np.sum((x - y) ** 2)))
            dist = np.array(dist)
            f.write(str(u.trajectory.time) + ',' + str(len(dist[dist <= threshold])) + '\n')
            print('Finished:', '{:.2f}'.format(ts.frame / len(u.trajectory) * 100), '%')


if __name__ == '__main__':
    u = mda.Universe('test/contacts/FcFF_SMP_r1_md.pdb', 'test/contacts/FcFF_SMP_r1_md_dt200.xtc')
    # cpd1 = get_aminoacid_sidechain_ring_atoms(u, 'cpd1')
    # cpd2 = get_aminoacid_sidechain_ring_atoms(u, 'cpd2')
    # phe = get_aminoacid_sidechain_ring_atoms(u, 'phe')
    # calc_intra_hydrophobic_contacts(u, 'cpd1', 'test/contacts/contacts_cpd1_cpd1_r0.5.csv')
    # calc_inter_hydrophobic_contacts(u, u, 'cpd1', 'cpd2', 'test/contacts/contacts_cpd1_cpd2_r0.5.csv')
    # calc_inter_hydrophobic_contacts(u, u, 'cpd1', 'phe', 'test/contacts/contacts_cpd1_phe_r0.5.csv')
    # calc_intra_hydrophobic_contacts(u, 'cpd2', 'test/contacts/contacts_cpd2_cpd2_r0.5.csv')
    # calc_inter_hydrophobic_contacts(u, u, 'cpd2', 'phe', 'test/contacts/contacts_cpd2_phe_r0.5.csv')
    # calc_intra_hydrophobic_contacts(u, 'phe', 'test/contacts/contacts_phe_phe_r0.5.csv')
    calc_inter_contacts(u, u, 'resid 3 and name O*', 'resname SMP and name N*', 'test/contacts/salt_bridge_r1.csv')
    # calc_intra_hydrophobic_contacts(u, 'phe', 'test/contacts/contacts_CPD-FF.csv')
